import copy
import logging
import signal
import threading
import time
from concurrent.futures.thread import ThreadPoolExecutor
from typing import Dict, List
from .job_store.interface import JobStore

from .data_struct import BucketDataStruct, Job, TimingWheelDataStruct
from .delay_queue.delay_queue import DelayQueue
from .retry_policy.interface import AddOrRunTimer, policy_pool
from .handler_interface import handler_pool


class BucketStore:
    """时间格对象存储"""
    lock: threading.Lock
    buckets: Dict[int, BucketDataStruct]
    bucket_idx: int

    def __init__(self):
        self.lock = threading.Lock()
        self.buckets = {}
        self.bucket_idx = None

    def new(self, current_time: int, wheel_id: int) -> BucketDataStruct:
        with self.lock:
            if self.bucket_idx is None:
                self.bucket_idx = 0
            else:
                self.bucket_idx += 1
            bucket = BucketDataStruct(current_time, [], self.bucket_idx, wheel_id)
            self.buckets[self.bucket_idx] = bucket
            return bucket

    def add_timer(self, timer: Job, bucket_idx: int):
        with self.buckets[bucket_idx].lock:
            self.buckets[bucket_idx].timers.insert(0, timer)

    def pop_timers(self, bucket_idx: int) -> List[Job]:
        """清空当前时间格的任务列表"""
        timers = []
        with self.buckets[bucket_idx].lock:
            self._set_expiration(bucket_idx, -1)
            while True:
                if len(self.buckets[bucket_idx].timers) > 0:
                    timers.append(self.buckets[bucket_idx].timers.pop())
                else:
                    break
        return timers

    def get_bucket(self, bucket_idx: int) -> BucketDataStruct:
        return self.buckets.get(bucket_idx)

    def _set_expiration(self, bucket_idx: int, expiration: int):
        if self.buckets[bucket_idx].expiration == expiration:
            return False
        else:
            self.buckets[bucket_idx].expiration = expiration
            return True

    def set_expiration(self, bucket_idx: int, expiration: int) -> bool:
        with self.buckets[bucket_idx].lock:
            return self._set_expiration(bucket_idx, expiration)

    def get_expiration(self, bucket_idx: int) -> int:
        with self.lock:
            return self.buckets[bucket_idx].expiration


class TimingWheelStore:
    """时间轮对象存储"""
    lock: threading.Lock
    wheels: Dict[int, TimingWheelDataStruct]

    def __init__(self):
        self.lock = threading.Lock()
        self.wheels = {}

    def add(self, wheel: TimingWheelDataStruct) -> bool:
        with self.lock:
            if wheel.interval in self.wheels:
                return False
            else:
                self.wheels[wheel.interval] = wheel
                return True

    def get(self, wheel_id: int) -> TimingWheelDataStruct:
        """获取时间轮对象"""
        with self.lock:
            return self.wheels.get(wheel_id)

    def get_current_time(self, wheel_id: int) -> int:
        with self.wheels[wheel_id].lock:
            return self.wheels.get(wheel_id).current_time

    def update_current_time(self, wheel_id: int, current_time: int):
        with self.wheels[wheel_id].lock:
            self.wheels[wheel_id].current_time = current_time


class TimingWheel(AddOrRunTimer):
    """时间轮维护类"""
    wheel_store: TimingWheelStore
    bucket_store: BucketStore
    executor: ThreadPoolExecutor
    quit_event = threading.Event
    logger: logging.Logger = None
    running = 0
    job_store: JobStore

    def __init__(self,
                 first_wheel_tick_ms: int,
                 wheel_size: int,
                 max_workers: int,
                 job_store: JobStore,
                 logger = None) -> None:
        """
        :param int first_wheel_tick_ms: 最底层时间轮时间格大小, 单位毫秒
        :param int wheel_size: 时间轮格子总数
        :param logging.Logger logger: 日志记录器
        """
        self.first_wheel_tick_ms = first_wheel_tick_ms
        self.wheel_size = wheel_size
        self.first_wheel_id = first_wheel_tick_ms * wheel_size
        if logger:
            self.logger = logger
        else:
            self.logger = logging.getLogger(self.__class__.__name__)
            self.logger.setLevel("DEBUG")
            fmt_str = "%(asctime)s %(name)s %(levelname)s %(module)s:%(funcName)s %(message)s"
            fmt = logging.Formatter(fmt_str)
            hd = logging.StreamHandler()
            hd.setLevel("DEBUG")
            hd.setFormatter(fmt)
            self.logger.addHandler(hd)
        self.bucket_dq = DelayQueue(wheel_size)
        self.wheel_store = TimingWheelStore()
        self.bucket_store = BucketStore()
        self.job_store = job_store
        self.executor = ThreadPoolExecutor(max_workers=max_workers)
        self.quit_event = threading.Event()

    def new_timingwheel(self,
                        tick_ms: int,
                        wheel_size: int,
                        start_ms: int) -> TimingWheelDataStruct:
        """
        :param int tick_ms: 时间轮时间格大小,单位毫秒
        :param int wheel_size: 时间轮格子总数
        :param int start_ms: 时间轮current_time值
        """
        if (tick_ms / 1000) < 0.001:
            raise ValueError("tick_ms 必须大于等于1ms")
        interval = tick_ms * wheel_size
        overflow_wheel_id = interval * wheel_size
        current_time = start_ms
        wheel = self.wheel_store.get(interval)
        if not wheel:
            buckets = []
            for _ in range(wheel_size):
                bucket = self.bucket_store.new(current_time, interval)
                buckets.append(bucket.bucket_idx)
            wheel = TimingWheelDataStruct(
                tick_ms,
                wheel_size,
                interval,
                current_time,
                buckets,
                overflow_wheel_id
            )
            self.wheel_store.add(wheel)
        return wheel

    def truncate(self, timestamp: int, tick: int) -> int:
        """
        将时间戳裁剪为tick的整倍数
        """
        if tick <= 0:
            return timestamp
        return timestamp - timestamp % tick

    def current_timestamp(self) -> int:
        """获取当前时间, 精度为毫秒"""
        return int(time.time() * 1000)

    def start(self):
        if self.running == 1:
            return
        signal.signal(signal.SIGQUIT, self.stop)
        signal.signal(signal.SIGINT, self.stop)
        start_ms = self.truncate(self.current_timestamp(), self.first_wheel_tick_ms)
        self.new_timingwheel(
            tick_ms=self.first_wheel_tick_ms,
            wheel_size=self.wheel_size,
            start_ms=start_ms
        )
        threading.Thread(target=self.start_delay_queue, daemon=True).start()
        threading.Thread(target=self.handle_timeout_job, daemon=True).start()
        self.running = 1

    def join(self):
        self.start()
        signal.pause()
        sleep_seconds = 1
        self.logger.warning(f"\n任务{sleep_seconds}秒后退出\n")
        time.sleep(sleep_seconds)

    def start_delay_queue(self):
        """启动延迟队列"""
        self.bucket_dq.poll(self.quit_event)

    def handle_timeout_job(self):
        while not self.quit_event.is_set():
            try:
                bucket_idx = self.bucket_dq.get()
                if self.quit_event.is_set():
                    self.logger.debug("程序已退出")
                    break
                bucket = self.bucket_store.get_bucket(bucket_idx)
                self.advance_clock(
                    bucket.wheel_id,
                    self.bucket_store.get_expiration(bucket_idx)
                )
                # 从到期的bucket中去除所有到期的任务提交到线程池执行
                timers = self.bucket_store.pop_timers(bucket_idx)
                for timer in timers:
                    self.add_or_run(timer)
            except Exception as err:
                print(err)

    def submit_timer(self, timer: Job):
        if self.quit_event.is_set():
            return
        self.executor.submit(self.timeout_func, timer)

    def stop(self, *args):
        self.quit_event.set()
        self.bucket_dq.wakup()
        self.executor.shutdown(wait=True)

    def timeout_func(self, timer: Job):
        """到期任务处理"""
        retry_policy = policy_pool.get(timer.retry_policy_name)
        @retry_policy.retry(self)
        def inner(timer: Job):
            handler = handler_pool.get(timer.handler_name)
            if handler is None:
                return
            handler.handle(**timer.func_kwargs)
        return inner(timer)

    def advance_clock(self, wheel_id: int, expiration: int):
        """推进时间轮当前时间"""
        current_time = self.wheel_store.get_current_time(wheel_id)
        w = self.wheel_store.get(wheel_id)
        if expiration >= current_time + w.tick:
            current_time = self.truncate(expiration, w.tick)
            self.wheel_store.update_current_time(wheel_id, current_time)
            overflow_wheel = self.wheel_store.get(w.overflow_wheel_id)
            if overflow_wheel:
                self.advance_clock(w.overflow_wheel_id, current_time)

    def get_or_create_timingwheel(self, expiration: int, wheel_id: int) -> TimingWheelDataStruct:
        """递归寻找或创建合适的时间轮"""
        tw: TimingWheelDataStruct = self.wheel_store.get(wheel_id)
        if not tw:
            if wheel_id == self.first_wheel_id:
                tw = self.new_timingwheel(
                    self.first_wheel_tick_ms,
                    self.wheel_size,
                    self.current_timestamp()
                )
            else:
                raise ValueError("时间轮不存在")
        current_time = copy.copy(tw.current_time)
        if expiration < current_time + tw.tick:
            # 任务已经过期, 不需要任何时间轮
            return None
        elif expiration < current_time + tw.interval:
            return tw
        else:
            # 递归向上层时间轮寻找位置存放任务
            overflow_wheel = self.wheel_store.get(tw.overflow_wheel_id)
            if not overflow_wheel:
                overflow_wheel: TimingWheelDataStruct = self.new_timingwheel(
                    tw.interval,
                    tw.wheel_size,
                    tw.current_time
                )
            return self.get_or_create_timingwheel(expiration, tw.overflow_wheel_id)

    def get_bucket_id(self, expiration: int, wheel_id: int) -> int:
        """计算出时间轮存放任务的格子(bucket)"""
        tw: TimingWheelDataStruct = self.wheel_store.get(wheel_id)
        virtual_id = int(expiration / tw.tick)
        return tw.bucket_idxs[virtual_id % tw.wheel_size]

    def add_timer(self, timer: Job) -> int:
        """
        添加任务到恰当的时间轮格子
        1. 找出合适的时间轮
        2. 找出合适的时间轮格子
        3. 存放任务到时间轮格子
        """
        if len(str(timer.expiration)) != 13:
            raise ValueError("Job.expiration 单位为毫秒")
        tw: TimingWheelDataStruct = self.get_or_create_timingwheel(
            timer.expiration,
            self.first_wheel_id
        )
        if tw is None:
            return -1
        bucket_idx = self.get_bucket_id(timer.expiration, tw.interval)
        self.bucket_store.add_timer(timer, bucket_idx)
        virtual_id = int(timer.expiration / tw.tick)
        if self.bucket_store.set_expiration(bucket_idx, virtual_id * tw.tick):
            bucket_expiration = self.bucket_store.get_expiration(bucket_idx)
            self.bucket_dq.offer(bucket_idx, bucket_expiration)
        return bucket_idx

    def add_or_run(self, timer: Job):
        """添加或运行任务"""
        result = self.add_timer(timer)
        if result == -1:
            self.submit_timer(timer)

    def new_job(self, timer: Job):
        self.job_store.add(timer)
        self.add_or_run(timer)